import math

import numpy as np

# import map_decrete as md
from enum import Enum




# 输入输出
"""
A star planning algorithm
input: map grid with 0 and 1, represents barrier and passable brick
        start point and end point
output: path from start to end
        the number of visited times of each node

"""

# 原始A*算法伪代码
"""
fake code:
function AStar(start, goal)
    openList := {start}
    closedList := {}

    while openList is not empty
        current := node in openList with the lowest f_cost

        if current is goal
            return path

        move current from openList to closedList

        for each neighbor of current
            if neighbor is not traversable or neighbor is in closedList
                skip to the next neighbor

            if neighbor is not in openList
                add neighbor to openList
                set neighbor's parent to current
                calculate neighbor's g_cost, h_cost, and f_cost

            else if neighbor is in openList and current's g_cost is lower than neighbor's g_cost
                update neighbor's parent and g_cost

    return no path found

"""


class Node:
    def __init__(self, x, y, traversable=True):
        # X代表行， y代表列
        self.x = x
        self.y = y
        self.traversable = traversable  # 0 表示障碍物，不可通行
        self.g_cost = 0
        self.h_cost = 0
        self.f_cost = 0
        self.parent = [0, 0]  # 使用列表存储父节点坐标，可以修改
        self.visited_times = 0


def create_node_matrix(map_grid):
    height, width = map_grid.shape
    nodes = [[Node(x, y, map_grid[x][y]) for y in range(width)] for x in range(height)]
    return nodes


class AstarPlanner:
    G_RATIO = 0.9
    print_flag = False

    def __init__(self, nodes, start, end, map_matrix):
        self.node_matrix = nodes
        self.xmax = len(nodes) - 1
        self.ymax = len(nodes[0]) - 1
        self.start = start
        self.end = end
        self.height = len(nodes)
        self.width = len(nodes[0])
        self.open_dict = dict()  # 元素是键值对 (x, y) : f_cost   list is no hashable
        self.closed_dict = dict()  # 元素是键值对 (x, y) : f_cost
        self.path = []
        self.map = map_matrix
        self.processed_grid = []

    @staticmethod
    def calc_diff_g_cost( parent_node, this_node):
        dx = parent_node[0] - this_node[0]
        dy = parent_node[1] - this_node[1]

        if dx == 0 and dy == 0:
            raise ValueError("The two nodes are the same")
        elif dx == 0 or dy == 0:
            return 1.0 * AstarPlanner.G_RATIO
        else:
            return 1.414 * AstarPlanner.G_RATIO

    # g和h的算法应该一致，且正确反映路程代价，直线一格为1，斜线一格为1.414
    def calc_h_cost(self, pos, h_cost):
        if not h_cost:
            dx = abs(pos[0] - self.end[0])
            dy = abs(pos[1] - self.end[1])
            d_min = min(dx, dy)
            d_max = max(dx, dy)
            return 1.414 * d_min + (d_max - d_min)
        return h_cost

    def add_to_open_set(self, pos_this, pos_parent):
        this_node = self.node_matrix[pos_this[0]][pos_this[1]]
        this_node.parent = pos_parent
        parent_node = self.node_matrix[pos_parent[0]][pos_parent[1]]
        this_node.h_cost = self.calc_h_cost(pos_this, this_node.h_cost)
        diff_g_cost = self.calc_diff_g_cost(pos_parent, pos_this)
        this_node.g_cost = parent_node.g_cost + diff_g_cost

        this_node.f_cost = this_node.g_cost + this_node.h_cost
        self.open_dict.update({(pos_this[0], pos_this[1]): this_node.f_cost})

    def get_path(self):
        path = [self.end]
        while path[-1] != self.start:
            parent = self.node_matrix[path[-1][0]][path[-1][1]].parent
            path.append(tuple(parent))
        return path

    def get_neighbours(self, pos):
        x, y = pos
        if 0 < x < self.xmax - 1:
            x_range = [x - 1, x, x + 1]
        elif x == 0:
            x_range = [0, 1]
        else:
            x_range = [self.xmax - 2, self.xmax - 1]

        if 0 < y < self.ymax - 1:
            y_range = [y - 1, y, y + 1]
        elif y == 0:
            y_range = [0, 1]
        else:
            y_range = [self.ymax - 2, self.ymax - 1]

        neighbours = [(i, j) for i in x_range for j in y_range if (i, j) != (x, y)]
        return neighbours

    def planning(self):
        # 判断初始化节点是否可行
        if self.map[self.start[0]][self.start[1]] == 0 or self.map[self.end[0]][self.end[1]] == 0:
            raise ValueError("The start node or the end node is not traversable")

        start_node = self.node_matrix[self.start[0]][self.start[1]]
        start_node.parent = self.start
        start_node.g_cost = 0
        start_node.h_cost = self.calc_h_cost(self.start, 0)
        start_node.f_cost = start_node.g_cost + start_node.h_cost
        self.open_dict.update({self.start: self.node_matrix[self.start[0]][self.start[1]].f_cost})

        while self.open_dict:
            current_pos = min(self.open_dict, key=self.open_dict.get)
            self.processed_grid.append(current_pos)
            current_node = self.node_matrix[current_pos[0]][current_pos[1]]
            # 打印当前节点信息，如果有必要
            if self.print_flag:
                print("current pos is {}， {}".format(current_pos[0], current_pos[1]))
                print("parent pos is {}， {}".format(current_node.parent[0], current_node.parent[1]))
                print("the f value is {}, the g value is {} and the h cost is {}"
                      .format(current_node.f_cost, current_node.g_cost, current_node.h_cost))
                # print("current direction is ", current_node.direction)
                print("~~~~\n")

            if current_pos == self.end:
                path = self.get_path()
                return path


            else:
                self.closed_dict.update({current_pos: current_node.f_cost})
                self.open_dict.pop(current_pos)

                # 更新parent (current) 的方向,如果指向障碍物，则置为 任意方向
                neighbour_list = self.get_neighbours(current_pos)

                for neighbour in neighbour_list:
                    neighbour_node = self.node_matrix[neighbour[0]][neighbour[1]]

                    if not neighbour_node.traversable or neighbour in self.closed_dict:
                        continue

                    if neighbour not in self.open_dict:
                        self.add_to_open_set(neighbour, current_pos)

                    elif neighbour in self.open_dict:
                        # 计算是否要更新parent
                        # 从current到neighbour这条路径的现实代价 在最理想的情况下都 劣于nbr的原路径，nbr就不必变换了
                        if (current_node.g_cost + 1 * AstarPlanner.G_RATIO) > neighbour_node.g_cost:
                            continue
                        # 从current到neighbour这条路径的现实代价 在最差的情况下都 优于nbr的原路径，nbr需要变换父节点
                        elif (current_node.g_cost + 2.414 * AstarPlanner.G_RATIO) < neighbour_node.g_cost:
                            self.add_to_open_set(neighbour, current_pos)
                        # 具体的计算代价，决定是否更新parent
                        else:
                            diff_g_cost = self.calc_diff_g_cost(current_pos, neighbour)
                            if current_node.g_cost + diff_g_cost < neighbour_node.g_cost:
                                self.node_matrix[neighbour[0]][neighbour[1]].parent = current_pos
                                self.node_matrix[neighbour[0]][neighbour[1]].g_cost = current_node.g_cost + diff_g_cost
                                self.node_matrix[neighbour[0]][neighbour[1]].f_cost = \
                                    (self.node_matrix[neighbour[0]][neighbour[1]].g_cost
                                        + self.node_matrix[neighbour[0]][neighbour[1]].h_cost)
                                self.open_dict.update({neighbour: self.node_matrix[neighbour[0]][neighbour[1]].f_cost})

        return []  # no path found


def main():
    # new_map = md.map_resize(md.map_grid, 50, 50)
    # np.save('npy/resized_map.npy', new_map)
    new_map = np.load("npy/resized_map_400.npy")
    create_node_matrix(new_map)
    start, end = (0, 0), (200, 300)
    astar = AstarPlanner(create_node_matrix(new_map), start, end, new_map)
    path = astar.planning()
    print(path)
    np.save('npy/raw_path_400.npy', path)


if __name__ == '__main__':
    main()
